Skip to content

PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction#41048

Open
copybara-service[bot] wants to merge 1 commit intomainfrom
test_900707821
Open

PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction#41048
copybara-service[bot] wants to merge 1 commit intomainfrom
test_900707821

Conversation

@copybara-service
Copy link
Copy Markdown

@copybara-service copybara-service bot commented Apr 16, 2026

PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction

Imported from GitHub PR #40920

  • New PropagateCallMetadata HLO pass that propagates metadata (op_name prefix and stack_frame_id) from kCall instructions into their called computations, recursing through nested control-flow (while, conditional) but not into embedded computations (reduce's to_apply, etc.)
  • Removes metadata propagation from CallInliner — the RecursivelyUpdateMetadata helper and per-instruction metadata update during inlining are no longer needed since the new pass handles this as a separate concern
  • Wired into GPU compiler pre-SPMD pipeline before CallInliner, so metadata is propagated while kCall ops still exist (including non-inlinable calls)
  • Test coverage for op_name propagation, stack frame concatenation, overflow protection, redundant prefix detection, nested calls, and idempotency

Motivation

The CallInliner only updates metadata for calls it actually inlines. Non-inlinable calls (e.g. calls with inlineable="false") were skipped entirely, leaving their callee instructions with incomplete metadata context. Extracting this into a standalone pass ensures all calls get metadata propagation regardless of inlining decisions.
Copybara import of the project:

--
3315b79 by Eugene Zhulenev ezhulenev@openxla.org:

[xla] Add a separate pass to propagate metadata across kCall instructions

Merging this change closes #40920

FUTURE_COPYBARA_INTEGRATE_REVIEW=#40920 from ezhulenev:propagate-metadata-pass 3315b79

@copybara-service copybara-service bot force-pushed the test_900707821 branch 9 times, most recently from bfcc606 to acd48c2 Compare April 17, 2026 06:57
…ll instruction

Imported from GitHub PR #40920

- New `PropagateCallMetadata` HLO pass that propagates metadata (op_name prefix and stack_frame_id) from `kCall` instructions into their called computations, recursing through nested control-flow (while, conditional) but not into embedded computations (reduce's to_apply, etc.)
- Removes metadata propagation from `CallInliner` — the `RecursivelyUpdateMetadata` helper and per-instruction metadata update during inlining are no longer needed since the new pass handles this as a separate concern
- Wired into GPU compiler pre-SPMD pipeline before `CallInliner`, so metadata is propagated while `kCall` ops still exist (including non-inlinable calls)
- Test coverage for op_name propagation, stack frame concatenation, overflow protection, redundant prefix detection, nested calls, and idempotency

#### Motivation

The `CallInliner` only updates metadata for calls it actually inlines. Non-inlinable calls (e.g. calls with inlineable="false") were skipped entirely, leaving their callee instructions with incomplete metadata context. Extracting this into a standalone pass ensures all calls get metadata propagation regardless of inlining decisions.
Copybara import of the project:

--
3315b79 by Eugene Zhulenev <ezhulenev@openxla.org>:

[xla] Add a separate pass to propagate metadata across kCall instructions

Merging this change closes #40920

FUTURE_COPYBARA_INTEGRATE_REVIEW=#40920 from ezhulenev:propagate-metadata-pass 3315b79
PiperOrigin-RevId: 900707821
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant